import world
import utils
from world import cprint
import torch
from datetime import datetime
import numpy as np
from tensorboardX import SummaryWriter
import time
import Procedure
import os
from  parse import  args
from os.path import join
import matplotlib.pyplot as plt
# ==============================
utils.set_seed(world.seed)
print(">>SEED:", world.seed)
# ==============================
import register
from register import dataset
from model import VAE

Recmodel = register.MODELS[world.model_name](world.config, dataset)
Recmodel = Recmodel.to(world.device)
Doc_vae = VAE(world.config['latent_dim_rec'],world.config['latent_dim_rec']//2,world.config['latent_dim_vae'],dataset.cat_type)
Doc_vae = Doc_vae.to(world.device)
Doc_vae.apply(utils.initialize_weights)
bpr = utils.BPRLoss(Recmodel, world.config,Doc_vae)

weight_file = utils.getFileName()
print(f"load and save to {weight_file}")
if world.LOAD:
    try:
        Recmodel.load_state_dict(torch.load(weight_file,map_location=torch.device('cpu')))
        world.cprint(f"loaded model weights from {weight_file}")
    except FileNotFoundError:
        print(f"{weight_file} not exists, start from beginning")
Neg_k = 1

# init tensorboard
if world.tensorboard:
    w : SummaryWriter = SummaryWriter(
                                    join(world.BOARD_PATH, time.strftime("%m-%d-%Hh%Mm%Ss-") + "-" + world.comment)
                                    )
else:
    w = None
    world.cprint("not enable tensorflowboard")


try:
    neg_epoch = 0
    best_result = None
    best_recall = 0
    item_emb = Recmodel.embedding_item.weight.detach().cpu().numpy()
    for epoch in range(world.TRAIN_epochs):
        start = time.time()
        if epoch %10 == 0 and epoch != 0:
            cprint("[TEST]")
            world.test_x.append(epoch)
            result = Procedure.Test(dataset, Recmodel, epoch, w, world.config['multicore'])
            if result['recall'][1]>best_recall:
                neg_epoch=0
                best_result = result
                best_recall = result['recall'][1]
                item_emb = Recmodel.embedding_item.weight.detach().cpu().numpy()
            elif neg_epoch< world.early_stop:
                neg_epoch+=1
            else:
                print("----early stop----")
                world.result_logger.logging(str(best_result))
                break
        output_information = Procedure.BPR_train_original(dataset, Recmodel, bpr, epoch, neg_k=Neg_k,w=w)
        world.loss_x.append(epoch)
        print(f'EPOCH[{epoch+1}/{world.TRAIN_epochs}] {output_information}')
        #torch.save(Recmodel.state_dict(), weight_file)
    time = "%s" % (datetime.now().strftime('%Y-%m-%d-%H-%M'))
    os.makedirs("./results/"+time, exist_ok=True)
    np.save("./results/" + time + "/item_emb.npy", item_emb)
    plt.figure()
    # plt.plot(world.loss_x, world.loss_y, color="red", label="Loss")
    plt.plot(world.test_x, world.tail, color="red", label="tail")
    plt.plot(world.test_x, world.tail_recall, color="purple", label="tail_recall")
    # plt.scatter(world.loss_x,world.loss_y,color='red')
    plt.legend()  # 无此语句会不显示右下角label
    plt.savefig("./results/" + time + '/loss_plot.png')
    plt.close()
    plt.figure()
    plt.plot(world.test_x, world.recall_y, color="green", label="Recall")
    plt.plot(world.test_x, world.ndcg_y, color="blue", label="NDCG")
    plt.legend()  # 无此语句会不显示右上角label
    plt.savefig("./results/" + time + '/Recall.png')
    plt.close()
    plt.figure()
    plt.plot(world.test_x, world.coverage, color="red", label="coverage")
    # plt.plot(world.test_x, world.hild, color="purple", label="MILD")
    plt.plot(world.test_x, world.tail_coverage, color="purple", label="tail_cov")
    plt.legend()
    plt.savefig("./results/" + time + '/Long_tail.png')
    plt.close()
    if world.is_show:
        plt.show()
finally:
    if world.tensorboard:
        w.close()